package flink.datastream;

import flink.model.Goods;
import flink.model.Order;
import org.apache.flink.streaming.api.collector.selector.OutputSelector;
import org.apache.flink.streaming.api.datastream.ConnectedStreams;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.datastream.SplitStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.CoMapFunction;
import org.junit.Test;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author binarylei
 * @version 2019-11-28
 */
public class OperatorDataStreamTest {

    @Test
    public void testUnion() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Order> ds1 = env.fromCollection(Arrays.asList(
                new Order(1L, "order1"),
                new Order(2L, "order2")));
        DataStreamSource<Order> ds2 = env.fromCollection(Arrays.asList(
                new Order(3L, "order3"),
                new Order(4L, "order4")));

        DataStream<Order> ds3 = ds1.union(ds2);
        ds3.print().setParallelism(1);
        env.execute();
    }

    @Test
    public void testConnect() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Order> ds1 = env.fromCollection(Arrays.asList(
                new Order(1L, "order1"),
                new Order(2L, "order2")));
        DataStreamSource<Goods> ds2 = env.fromCollection(Arrays.asList(
                new Goods(3L, "goods3"),
                new Goods(4L, "goods4")));

        ConnectedStreams<Order, Goods> connectedStreams = ds1.connect(ds2);

        SingleOutputStreamOperator<String> ds3 = connectedStreams.map(
                new CoMapFunction<Order, Goods, String>() {
            @Override
            public String map1(Order order) throws Exception {
                return order.name;
            }

            @Override
            public String map2(Goods goods) throws Exception {
                return goods.name;
            }
        });
        ds3.print().setParallelism(1);
        env.execute();
    }

    @Test
    public void testSplit() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<Order> ds1 = env.fromCollection(Arrays.asList(
                new Order(1L, "order1"),
                new Order(2L, "order2")));

        SplitStream<Order> splitStream = ds1.split(new OutputSelector<Order>() {
            @Override
            public Iterable<String> select(Order order) {
                List<String> out = new ArrayList<>();
                out.add(order.id % 2 == 0 ? "even" : "odd");
                return out;
            }
        });
        DataStream<Order> ds2 = splitStream.select("even", "odd");
        DataStream<Order> ds3 = splitStream.select("even");
        DataStream<Order> ds4 = splitStream.select("odd");

        ds3.print().setParallelism(1);
        env.execute();
    }

}
